"""
知识库检索器
基于ChromaDB实现RAG检索
"""
from typing import List, Dict, Any, Optional
from loguru import logger

from app.db.vector_db import vector_db_manager
from app.config import settings


class KnowledgeRetriever:
    """知识库检索器"""
    
    def __init__(self):
        """初始化检索器"""
        self.vector_db = vector_db_manager
        self.max_results = settings.MAX_SEARCH_RESULTS
        self.use_rerank = settings.USE_RERANK
        self.rerank_provider = settings.RERANK_PROVIDER
    
    def _rerank_with_dashscope(self, query: str, documents: List[str]) -> List[float]:
        """
        使用阿里云DashScope API进行Rerank
        
        Args:
            query: 查询文本
            documents: 文档列表
            
        Returns:
            相关性分数列表
        """
        try:
            import dashscope
            from dashscope import TextReRank
            
            if not settings.DASHSCOPE_API_KEY:
                logger.error("DASHSCOPE_API_KEY未配置")
                return [0.0] * len(documents)
            
            dashscope.api_key = settings.DASHSCOPE_API_KEY
            
            # 调用Rerank API
            response = TextReRank.call(
                model=settings.RERANK_MODEL,
                query=query,
                documents=documents,
                top_n=min(len(documents), settings.RERANK_TOP_N),
                return_documents=False
            )
            
            if response.status_code != 200:
                logger.error(f"Rerank API失败: {response.message}")
                return [0.0] * len(documents)
            
            # 提取分数
            scores = [0.0] * len(documents)
            for result in response.output.results:
                scores[result.index] = result.relevance_score
            
            return scores
            
        except Exception as e:
            logger.error(f"Rerank调用失败: {e}")
            return [0.0] * len(documents)
        
    async def search(
        self,
        query: str,
        k: int = None,
        threshold: float = 0.3,  # 临时降低以便调试
        filters: Optional[Dict] = None
    ) -> List[Dict[str, Any]]:
        """
        检索相关文档
        
        Args:
            query: 查询文本
            k: 返回结果数量
            threshold: 相关性阈值（0-1）
            filters: 元数据过滤条件
            
        Returns:
            检索结果列表
        """
        try:
            k = k or self.max_results
            
            logger.info(f"开始检索: {query[:50]}... | k={k} | threshold={threshold}")
            
            # 使用向量数据库检索
            results = self.vector_db.similarity_search(
                query=query,
                k=k,
                threshold=threshold
            )
            
            logger.info(f"检索完成，找到 {len(results)} 个相关文档")
            
            # 格式化结果
            formatted_results = []
            for result in results:
                formatted_results.append({
                    "content": result.get("content", ""),
                    "metadata": result.get("metadata", {}),
                    "relevance_score": result.get("relevance_score", 0),
                    "distance": result.get("distance", 1.0),
                })
            
            return formatted_results
            
        except Exception as e:
            logger.error(f"检索失败: {e}")
            return []
    
    async def search_with_rerank(
        self,
        query: str,
        k: int = None,
        threshold: float = 0.3,
        recall_multiplier: int = None
    ) -> List[Dict[str, Any]]:
        """
        两阶段检索：召回 + Rerank精排（远程API）
        
        Args:
            query: 查询文本
            k: 最终返回结果数
            threshold: 召回阈值
            recall_multiplier: 召回倍数
            
        Returns:
            精排后的Top-K结果
        """
        try:
            k = k or self.max_results
            recall_multiplier = recall_multiplier or settings.RERANK_RECALL_MULTIPLIER
            
            # ========== 第1阶段：召回 ==========
            recall_k = min(k * recall_multiplier, 100)
            logger.info(f"[召回] 查询: {query[:50]}... | 召回: {recall_k}个")
            
            candidates = await self.search(query, k=recall_k, threshold=threshold)
            
            if not candidates:
                return []
            
            if len(candidates) <= k or len(candidates) < settings.RERANK_MIN_CANDIDATES:
                logger.info(f"候选数({len(candidates)})不足，跳过Rerank")
                return candidates[:k]
            
            logger.info(f"[召回] 成功召回 {len(candidates)} 个")
            
            if not self.use_rerank:
                return candidates[:k]
            
            # ========== 第2阶段：Rerank精排 ==========
            logger.info(f"[精排] 使用 {settings.RERANK_MODEL}")
            
            documents = [doc['content'] for doc in candidates]
            
            if self.rerank_provider == "dashscope":
                rerank_scores = self._rerank_with_dashscope(query, documents)
            else:
                logger.error(f"不支持的提供商: {self.rerank_provider}")
                return candidates[:k]
            
            # 更新分数
            for doc, score in zip(candidates, rerank_scores):
                doc['original_score'] = doc.get('relevance_score', 0)
                doc['rerank_score'] = float(score)
                doc['relevance_score'] = float(score)
            
            # 重排序
            candidates.sort(key=lambda x: x['rerank_score'], reverse=True)
            final_results = candidates[:k]
            
            if final_results:
                logger.info(
                    f"[精排] 完成 | Top1: "
                    f"原始={final_results[0].get('original_score', 0):.3f}, "
                    f"精排={final_results[0]['rerank_score']:.3f}"
                )
            
            return final_results
            
        except Exception as e:
            logger.error(f"Rerank失败: {e}，回退到普通检索")
            return await self.search(query, k or self.max_results, threshold)
    
    def format_context(
        self,
        documents: List[Dict[str, Any]],
        max_length: int = 2000
    ) -> str:
        """
        格式化检索到的文档为上下文
        
        Args:
            documents: 文档列表
            max_length: 最大长度（字符数）
            
        Returns:
            格式化的上下文文本
        """
        if not documents:
            return ""
        
        context_parts = []
        total_length = 0
        
        for i, doc in enumerate(documents, 1):
            content = doc.get("content", "")
            score = doc.get("relevance_score", 0)
            
            # 添加文档标识和相关性
            doc_text = f"[相关文档 {i}]（相关度：{score:.2f}）\n{content}\n"
            
            # 检查是否超过最大长度
            if total_length + len(doc_text) > max_length:
                break
            
            context_parts.append(doc_text)
            total_length += len(doc_text)
        
        return "\n".join(context_parts)
    
    def extract_keywords(self, query: str) -> List[str]:
        """
        从查询中提取关键词（简单实现）
        
        Args:
            query: 查询文本
            
        Returns:
            关键词列表
        """
        # 简单的关键词提取
        # TODO: 可以使用更复杂的NLP方法
        import re
        
        # 移除标点符号
        cleaned = re.sub(r'[^\w\s]', ' ', query)
        
        # 分词（简单按空格分）
        words = cleaned.split()
        
        # 过滤停用词（简化版）
        stopwords = {'的', '了', '是', '在', '我', '有', '和', '就', '不', '人', '都', '一', '一个'}
        keywords = [w for w in words if w not in stopwords and len(w) > 1]
        
        return keywords[:5]  # 最多返回5个关键词


# 全局检索器实例
knowledge_retriever = KnowledgeRetriever()


# 便捷函数
async def search_knowledge(
    query: str,
    k: int = None,
    threshold: float = 0.6
) -> List[Dict[str, Any]]:
    """便捷的知识检索函数"""
    return await knowledge_retriever.search(query, k, threshold)

